# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     datalist
   Description :   
   Author :       lth
   date：          2022/10/14
-------------------------------------------------
   Change Activity:
                   2022/10/14 3:08: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import os

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from torchvision.datasets import CIFAR10
#
# dataset = CIFAR10(
#     root="./data", train=True, download=True,
#     transform=transforms.Compose([
#         transforms.RandomHorizontalFlip(),
#         transforms.ToTensor(),
#         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
#     ]
#     )
# )

train_trasform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

test_trasform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])


class CelebAHQ(Dataset):
    def __init__(self, path="E:\Datasets2\CelebAMask-HQ\CelebA-HQ-img", mode="train"):
        super(CelebAHQ, self).__init__()
        self.data = self.get_data(path)
        self.image_size = (128,128)
        self.mode = mode
    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        path = self.data[index]
        img = Image.open(path).convert("RGB")
        img = img.resize(self.image_size)

        if self.mode == "train":
            return train_trasform(img)
        else:
            return test_trasform(img)

    @staticmethod
    def get_data(path):
        data = []
        for root, dir, files in os.walk(path):
            for file in files:
                data.append(os.path.join(root, file))
        return data
